from datetime import datetime
from zoneinfo import ZoneInfo
import os
import base64
from typing import Tuple


from gmutil import SM4Encryptor, sm3_hash, SM4Decryptor
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.padding import PKCS7
from cryptography.hazmat.primitives.hashes import SHA256, Hash

VERSION = 0x01
ALGORITHM_SM4_SM3 = 0x01
ALGORITHM_AES_SHA256 = 0x02


def encrypt01(secret_key: bytes):
    def _encrypt(data):
        enc = SM4Encryptor(secret_key, 'ECB', 'PKCS7')
        return enc.update(data) + enc.finalize()
    return _encrypt

def decrypt01(secret_key: bytes):
    def _decrypt(data):
        dec = SM4Decryptor(secret_key, 'ECB', 'PKCS7')
        return dec.update(data) + dec.finalize()
    return _decrypt

def hash01(data):
    return sm3_hash(data)

def encrypt02(secret_key):
    def _encrypt(data):
        cipher = Cipher(algorithm=algorithms.AES(secret_key), mode=modes.ECB())
        enc = cipher.encryptor()
        padder = PKCS7(128).padder()
        return enc.update(padder.update(data) + padder.finalize()) + enc.finalize()
    return _encrypt

def decrypt02(secret_key):
    def _decrypt(data):
        cipher = Cipher(algorithm=algorithms.AES(secret_key), mode=modes.ECB())
        dec = cipher.decryptor()
        unpadder = PKCS7(128).unpadder()
        return unpadder.update(dec.update(data) + dec.finalize()) + unpadder.finalize()
    return _decrypt


def hash02(data):
    sha = Hash(SHA256())
    sha.update(data)
    return sha.finalize()


def cucc_encode(university_id: bytes, user_id: bytes, algorithm: int, secret_key: bytes) -> str:
    if len(university_id) != 5:
        raise ValueError()

    buffer = bytearray()  # 输出缓存
    to_encrypt = bytearray()  # 加密缓存
    to_hash = bytearray()  # 哈希缓存

    buffer.append(VERSION)  # 版本号
    buffer.extend(university_id)  # 学校标识码
    now = datetime.now(ZoneInfo('Asia/Shanghai'))
    timestamp = bytes.fromhex(now.strftime('%Y%m%d%H%M%S'))
    # print(now.strftime('%Y-%m-%d %H:%M:%S'))

    buffer.extend(timestamp)  # 时间戳
    to_hash.extend(timestamp)

    to_encrypt.extend(user_id)
    to_hash.extend(user_id)

    nonce = os.urandom(4)
    to_encrypt.extend(nonce)
    to_hash.extend(nonce)

    if algorithm == ALGORITHM_SM4_SM3:
        encrypted = encrypt01(secret_key)(to_encrypt)
        hashed = hash01(to_hash)
    elif algorithm == ALGORITHM_AES_SHA256:
        encrypted = encrypt02(secret_key)(to_encrypt)
        hashed = hash02(to_hash)
    else:
        raise ValueError()
    # print(buffer.hex())

    custom_len = 1 + len(encrypted) + len(hashed)
    buffer.append(custom_len)
    buffer.append(algorithm)
    buffer.extend(encrypted)
    buffer.extend(hashed)

    code = base64.urlsafe_b64encode(buffer).decode()
    # print(code)

    return code


def cucc_decode(code, secret_key: bytes) -> Tuple[str, str, str]:
    octets = memoryview(base64.urlsafe_b64decode(code))
    version = octets[0]
    if version != VERSION:
        raise ValueError('CUCC版本号错误')
    university_id = octets[1:6]
    timestamp = octets[6:13]

    custom_len = octets[13]
    if len(octets) != custom_len + 14:
        raise ValueError('赋码高校自定义内容长度错误')

    algorithm = octets[14]
    encrypted = memoryview(octets[15:-32])
    if algorithm == ALGORITHM_SM4_SM3:
        decrypted = decrypt01(secret_key)(encrypted)
        to_hash = bytearray()
        to_hash.extend(timestamp)
        to_hash.extend(decrypted)
        if hash01(to_hash) != octets[-32:]:
            raise ValueError('SM3哈希值校验错误')
        user_id = decrypted[0:-4].decode()
    elif algorithm == ALGORITHM_AES_SHA256:
        decrypted = decrypt02(secret_key)(encrypted)
        to_hash = bytearray()
        to_hash.extend(timestamp)
        to_hash.extend(decrypted)
        if hash02(to_hash) != octets[-32:]:
            raise ValueError('SHA256哈希值校验错误')
        user_id = decrypted[0:-4].decode()
    else:
        raise ValueError('无法识别的算法代码{}'.format(algorithm))

    return university_id.hex(), timestamp.hex(), user_id

